import os
import cv2
import json
import numpy as np

def convert_txt_to_labelme_json(root_dir, image_formats=('.png', '.jpg')):
    # YOLO txt to LabelMe JSON conversion
    labelList = ['phone','smoke']
    
    for subdir, _, files in os.walk(root_dir):
        for txt_file in files:
            if not txt_file.endswith('.txt'):
                continue
            
            txt_path = os.path.join(subdir, txt_file)
            labelme_json = {
                'version': '5.3.1',
                'flags': {},
                'shapes': [],
                'imagePath': None,
                'imageData': None,
                'imageHeight': None,
                'imageWidth': None,
            }

            # Find the corresponding image path (try multiple formats)
            image_name = None
            for fmt in image_formats:
                candidate_image_name = txt_file.rsplit(".", 1)[0] + fmt
                candidate_image_path = os.path.join(subdir, candidate_image_name)
                if os.path.exists(candidate_image_path):
                    image_name = candidate_image_name
                    image_path = candidate_image_path
                    break
            
            if image_name is None:
                print(f'Warning: No image found for {txt_file} with formats {image_formats}')
                continue

            # Load image to get dimensions
            image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
            if image is None:
                print(f"Error: Failed to load image {image_path}")
                continue
            h, w = image.shape[:2]

            labelme_json['imagePath'] = image_name
            labelme_json['imageHeight'] = h
            labelme_json['imageWidth'] = w
            
            # Read txt file content
            with open(txt_path, 'r') as t:
                lines = t.readlines()
                for line in lines:
                    content = line.strip().split()
                    
                    # Parse YOLO format content
                    try:
                        label = labelList[int(content[0])]
                        x_center, y_center, object_width, object_height = map(float, content[1:5])
                        score = float(content[5]) if len(content) > 5 else 0.99
                        
                        top_left_x = (x_center - object_width / 2) * w
                        top_left_y = (y_center - object_height / 2) * h
                        bottom_right_x = (x_center + object_width / 2) * w
                        bottom_right_y = (y_center + object_height / 2) * h

                        shape = {
                            'label': label,
                            'score': score,
                            'group_id': None,
                            'shape_type': 'rectangle',
                            'flags': {},
                            'points': [
                                [top_left_x, top_left_y],
                                [bottom_right_x, bottom_right_y]
                            ]
                        }
                        labelme_json['shapes'].append(shape)
                    except (IndexError, ValueError) as e:
                        print(f"Error in file {txt_path}: {e}")
                        continue

            # Create JSON file
            json_name = txt_file.rsplit('.', 1)[0] + '.json'
            json_path = os.path.join(subdir, json_name)

            with open(json_path, 'w') as fd:
                json.dump(labelme_json, fd, indent=4)
            
            print(f"Saved JSON: {json_path}")

if __name__ == "__main__":
    root_dir = 'smoke'  # The root directory containing images and txt files in subdirectories
    convert_txt_to_labelme_json(root_dir, image_formats=('.png', '.jpg'))
